"""
人员检测模块
使用YOLOv11进行实时人员检测
"""

import cv2
import numpy as np
from ultralytics import YOLO
import supervision as sv
import logging
from typing import Optional, List, Tuple

class PersonDetector:
    """YOLOv11人员检测器"""
    
    def __init__(self, model_path: str = "yolo11n.pt", 
                 confidence_threshold: float = 0.5,
                 iou_threshold: float = 0.45,
                 classes: List[int] = [0]):
        """
        初始化检测器
        
        Args:
            model_path: 模型文件路径
            confidence_threshold: 置信度阈值
            iou_threshold: IoU阈值
            classes: 检测类别列表，默认[0]表示只检测人员
        """
        self.model_path = model_path
        self.confidence_threshold = confidence_threshold
        self.iou_threshold = iou_threshold
        self.classes = classes
        
        # 设置日志
        self.logger = logging.getLogger(__name__)
        
        # 初始化模型
        self.model = None
        self._load_model()
        
    def _load_model(self):
        """加载YOLOv11模型"""
        try:
            self.model = YOLO(self.model_path)
            self.logger.info(f"成功加载模型: {self.model_path}")
        except Exception as e:
            self.logger.error(f"加载模型失败: {e}")
            raise
    
    def detect(self, frame: np.ndarray) -> sv.Detections:
        """
        检测人员
        
        Args:
            frame: 输入图像帧
            
        Returns:
            supervision.Detections: 检测结果
        """
        if self.model is None:
            self.logger.error("模型未加载")
            return sv.Detections.empty()
        
        try:
            # 使用YOLOv11进行检测
            results = self.model(
                frame,
                conf=self.confidence_threshold,
                iou=self.iou_threshold,
                classes=self.classes,
                verbose=False
            )
            
            # 转换为supervision格式
            detections = sv.Detections.from_ultralytics(results[0])
            
            return detections
            
        except Exception as e:
            self.logger.error(f"检测过程出错: {e}")
            return sv.Detections.empty()
    
    def get_model_info(self) -> dict:
        """获取模型信息"""
        if self.model is None:
            return {}
        
        return {
            "model_path": self.model_path,
            "confidence_threshold": self.confidence_threshold,
            "iou_threshold": self.iou_threshold,
            "classes": self.classes
        }
    
    def update_thresholds(self, confidence: float = None, iou: float = None):
        """更新检测阈值"""
        if confidence is not None:
            self.confidence_threshold = confidence
            self.logger.info(f"置信度阈值更新为: {confidence}")
        
        if iou is not None:
            self.iou_threshold = iou
            self.logger.info(f"IoU阈值更新为: {iou}")


class DetectionVisualizer:
    """检测结果可视化器"""
    
    def __init__(self):
        # 创建标注器
        self.box_annotator = sv.BoxAnnotator(
            thickness=2
        )
        
        self.label_annotator = sv.LabelAnnotator(
            text_scale=0.7,
            text_color=sv.Color.WHITE
        )
    
    def annotate_frame(self, frame: np.ndarray, detections: sv.Detections) -> np.ndarray:
        """
        在帧上标注检测结果
        
        Args:
            frame: 输入帧
            detections: 检测结果
            
        Returns:
            标注后的帧
        """
        # 创建标签
        labels = []
        if detections.confidence is not None:
            labels = [
                f"Person {confidence:.2f}"
                for confidence in detections.confidence
            ]
        else:
            labels = ["Person"] * len(detections)
        
        # 绘制边界框
        annotated_frame = self.box_annotator.annotate(
            scene=frame.copy(),
            detections=detections
        )
        
        # 绘制标签
        annotated_frame = self.label_annotator.annotate(
            scene=annotated_frame,
            detections=detections,
            labels=labels
        )
        
        return annotated_frame
    
    def draw_detection_info(self, frame: np.ndarray, detections: sv.Detections, 
                          position: Tuple[int, int] = (10, 30)) -> np.ndarray:
        """
        在帧上绘制检测信息
        
        Args:
            frame: 输入帧
            detections: 检测结果
            position: 文本位置
            
        Returns:
            绘制信息后的帧
        """
        x, y = position
        
        # 绘制检测到的人数
        person_count = len(detections)
        text = f"Detected: {person_count}"
        
        cv2.putText(
            frame, text, (x, y),
            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2
        )
        
        return frame


if __name__ == "__main__":
    # 测试代码
    import time
    
    # 配置日志
    logging.basicConfig(level=logging.INFO)
    
    # 创建检测器
    detector = PersonDetector()
    visualizer = DetectionVisualizer()
    
    # 测试图像检测
    cap = cv2.VideoCapture(0)  # 使用摄像头测试
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        # 检测
        start_time = time.time()
        detections = detector.detect(frame)
        end_time = time.time()
        
        # 可视化
        annotated_frame = visualizer.annotate_frame(frame, detections)
        annotated_frame = visualizer.draw_detection_info(annotated_frame, detections)
        
        # 显示FPS
        fps = 1 / (end_time - start_time)
        cv2.putText(
            annotated_frame, f"FPS: {fps:.1f}", (10, 70),
            cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2
        )
        
        cv2.imshow("Person Detection Test", annotated_frame)
        
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    
    cap.release()
    cv2.destroyAllWindows() 